package main

var (
	MaxDepth  = 5
	MaxObject = 10
)
// 四元树又称四叉树是一种树状数据结构，在每一个节点上会有四个子区块。四元树常应用于二维空间数据的分析与分类。
// 它将数据区分成为四个象限。数据范围可以是方形或矩形或其他任意形状。
type Rect struct {
	X, Y int
	W, H int
}

type QuadTree struct {
	Depth    int
	Bound    *Rect
	Children [4]*QuadTree
	Objects  []*Rect
	Leaf     bool
}

func NewQuadTree(depth int, bound *Rect) *QuadTree {
	return &QuadTree{
		Depth: depth,
		Bound: bound,
		Leaf:  true,
	}
}

// Clear
func (this *QuadTree) Clear() {
	if this.Objects != nil {
		this.Objects = this.Objects[:0]
	}
	this.Leaf = true
	for _, n := range this.Children {
		if n != nil {
			n.Clear()
		}
	}
}

// Split the children node bound
func (this *QuadTree) split() {
	if this.Children[0] == nil {
		depth := this.Depth + 1
		x, y := this.Bound.X, this.Bound.Y
		w, h := this.Bound.W/2, this.Bound.H/2
		this.Children[0] = NewQuadTree(depth, &Rect{x + w, y, w, h})
		this.Children[1] = NewQuadTree(depth, &Rect{x, y, w, h})
		this.Children[2] = NewQuadTree(depth, &Rect{x, y + h, w, h})
		this.Children[3] = NewQuadTree(depth, &Rect{x + w, y + h, w, h})
	}
	this.Leaf = false
}

// Determine which quadrant the object completely fit within
func (this *QuadTree) getQuadrant(rect *Rect) int {
	quadrant := 0
	xmid := this.Bound.X + this.Bound.W/2
	ymid := this.Bound.Y + this.Bound.H/2

	// can completely fit within the top quadrants
	top := (rect.Y < ymid && rect.Y+rect.H < ymid)

	// can completely fit within the bottom quadrants
	bottom := (rect.Y > ymid)

	if rect.X < xmid && rect.X+rect.W < xmid {
		// can completely fit within the left quadrants
		if top {
			quadrant = 2
		} else if bottom {
			quadrant = 3
		}
	} else if rect.X > xmid {
		// can completely fit within the right quadrants
		if top {
			quadrant = 1
		} else if bottom {
			quadrant = 4
		}
	}
	return quadrant
}

// Insert the object into the quadtree
func (this *QuadTree) Insert(rect *Rect) {
	if !this.Leaf {
		q := this.getQuadrant(rect)
		if q != 0 {
			this.Children[q-1].Insert(rect)
			return
		}
	}

	this.Objects = append(this.Objects, rect)
	if len(this.Objects) > MaxObject && this.Depth < MaxDepth {
		if this.Leaf {
			this.split()
		}
		for i := 0; i < len(this.Objects); {
			q := this.getQuadrant(this.Objects[i])
			if q != 0 {
				r := this.Objects[i]
				this.Objects[i] = this.Objects[len(this.Objects)-1]
				this.Objects = this.Objects[:len(this.Objects)-1]
				this.Children[q-1].Insert(r)
			} else {
				i++
			}
		}
	}
}

// Return all objects that could collide with the given object
func (this *QuadTree) Retrieve(rect *Rect) []*Rect {
	objs := append([]*Rect{}, this.Objects...)
	q := this.getQuadrant(rect)
	if !this.Leaf {
		if q != 0 {
			objs = append(objs, this.Children[q-1].Retrieve(rect)...)
		} else {
			for i := 0; i < len(this.Children); i++ {
				objs = append(objs, this.Children[i].Retrieve(rect)...)
			}
		}
	}
	return objs
}
